#include "Binding.h"
#include <unordered_map>
#include <map>
#include <set>
#include <vector>
#include <cmath>
#include <algorithm>

std::vector<BindingInfo> computeBindings(
    const std::vector<Node>& nodes,
    const std::unordered_map<int, std::vector<EdgeSegment>>& skeletons,
    const std::vector<std::vector<Edge>>& clusteredEdges,
    int k,
    double sigma
) {
    std::unordered_map<int, Node> nodeMap;
    for (const auto& node : nodes) {
        nodeMap[node.id] = node;
    }

    std::vector<BindingInfo> bindings;

    // 预处理每个聚类的骨架点（去重后扁平化成点集合）
    std::unordered_map<int, std::vector<std::pair<double, double>>> skeletonPointsByCluster;
    for (const auto& [clusterId, segments] : skeletons) {
        std::set<std::pair<double, double>> uniquePoints;
        for (const auto& segment : segments) {
            uniquePoints.insert(segment.first);
            uniquePoints.insert(segment.second);
        }
        skeletonPointsByCluster[clusterId] = std::vector<std::pair<double, double>>(uniquePoints.begin(), uniquePoints.end());
    }

    // 主体逻辑：遍历每个聚类的边
    for (size_t clusterId = 0; clusterId < clusteredEdges.size(); ++clusterId) {
        const auto& edges = clusteredEdges[clusterId];
        const auto& skeletonPoints = skeletonPointsByCluster[clusterId];

        for (const auto& edge : edges) {
            const auto& p1 = nodeMap[edge.from];
            const auto& p2 = nodeMap[edge.to];

            std::pair<double, double> midpoint = {
                (p1.lon + p2.lon) / 2.0,
                (p1.lat + p2.lat) / 2.0
            };

            // 如果当前聚类没有骨架点，直接绑定到自身中点
            if (skeletonPoints.empty()) {
                bindings.push_back({ edge, midpoint, midpoint, static_cast<int>(clusterId) });
                continue;
            }

            // 计算与所有骨架点的距离平方
            std::vector<std::pair<std::pair<double, double>, double>> distPoints;
            for (const auto& skelPt : skeletonPoints) {
                double dx = skelPt.first - midpoint.first;
                double dy = skelPt.second - midpoint.second;
                double distSq = dx * dx + dy * dy;
                distPoints.push_back({ skelPt, distSq });
            }

            // 取最近 k 个骨架点
            std::nth_element(distPoints.begin(), distPoints.begin() + std::min(k, static_cast<int>(distPoints.size())) - 1, distPoints.end(),
                [](const auto& a, const auto& b) { return a.second < b.second; });

            std::vector<std::pair<std::pair<double, double>, double>> nearest(
                distPoints.begin(),
                distPoints.begin() + std::min(k, static_cast<int>(distPoints.size()))
            );

            // 高斯加权平均
            double sumX = 0.0, sumY = 0.0, totalWeight = 0.0;
            for (const auto& [point, distSq] : nearest) {
                double weight = std::exp(-distSq / (2 * sigma * sigma));
                sumX += point.first * weight;
                sumY += point.second * weight;
                totalWeight += weight;
            }

            std::pair<double, double> weightedPoint = {
                sumX / totalWeight,
                sumY / totalWeight
            };

            bindings.push_back({ edge, midpoint, weightedPoint, static_cast<int>(clusterId) });
        }
    }

    return bindings;
}
